import numpy, os, array, cv2
import pathlib

# b2m算法实现
# 将二进制代码转换为灰度图
# 原始apk位置（apk解压后路径）
unzip_path = '..\\data\\repack-unzip\\'
# 数据保存父目录
databasepath = '..\\datasets\\'

unzip_list = os.listdir(unzip_path)
for odir in unzip_list:
    # 处理原dex文件
    dirpath = os.path.join(unzip_path, odir)
    origin_dex_path = dirpath + '\\classes.dex'
    of = open(origin_dex_path, 'rb')
    ln = os.path.getsize(origin_dex_path)
    width = 512
    rem = ln % width
    length = int(ln / width)
    arr = array.array("B")
    arr.fromfile(of, ln - rem)
    of.close()
    # 灰度图生成
    oimage = numpy.reshape(arr, (length, width))
    oimage = numpy.uint8(oimage)
    # 灰度图缩放
    # INTER_AREA方法插值计算缩放
    # https://zhuanlan.zhihu.com/p/38493205
    oimage = cv2.resize(oimage, (256, 256), interpolation=cv2.INTER_AREA)
    opath = databasepath + odir
    os.mkdir(opath)
    databasefn = opath + '\\' + odir + '.png'
    cv2.imwrite(databasefn, oimage)
    print("原dex文件灰度图生成完毕:" + odir + "\n"
          "继续生成重打包dex灰度图")

    # 重打包dex灰度图生成
    repack_unzip_dir_list = os.listdir(unzip_path + odir)
    for rdir in repack_unzip_dir_list:
        repack_dex_path = unzip_path + odir + '\\' + rdir + "\\classes.dex"
        rf = open(repack_dex_path, 'rb')
        ln = os.path.getsize(repack_dex_path)
        width = 512
        rem = ln % rem
        length = int(ln / width)
        arr = array.array("B")
        arr.fromfile(rf, ln - rem)
        rf.close()
        # 灰度图生成
        rimage = numpy.reshape(arr, (length, width))
        rimage = numpy.uint8(rimage)
        # 灰度图缩放
        # INTER_AREA方法插值计算缩放
        # https://zhuanlan.zhihu.com/p/38493205
        rimage = cv2.resize(rimage, (256, 256), interpolation=cv2.INTER_AREA)
        # 直接保存到对应源灰度图目录下
        cv2.imwrite(opath + '\\' + rdir + '.png', rimage)
        print("重打包dex文件灰度图生成完毕:" + rdir)
